MB = 1024 * 1024
import json

class StrategyEvaluator:
    def __init__(self, num_microbatches=32):
        self.num_microbatches = num_microbatches
        # 添加基础性能参数
        self.device_tflops = 100  # 单卡算力100 TFLOPS
        self.bandwidth = 200  # GB/s

    def load_strategy(self, file_path):
        """从JSON文件加载策略配置"""
        with open(file_path, 'r') as f:
            return json.load(f)

    def calculate_stage_time(self, stage):
        """更精确的阶段时间计算"""
        stage_time = 0
        for i, strategy in enumerate(stage["strategies"]):
            dp, tp = strategy["dp"], strategy["tp"]

            # 模拟不同层类型的计算量（关键改进）
            if "embedding" in stage.get("layer_types", [""])[i]:
                flops = 5e12  # Embedding层计算量更大
            else:
                flops = 1e12  # 普通层

            # 计算时间 = 计算量 / (设备算力 * 并行度)
            compute_time = (flops / 1e12) / (self.device_tflops * dp * tp)

            # 通信时间 = 参数量/(带宽*TP) + 梯度量/(带宽*DP)
            param_size = 1e9 / (tp * 1e9)  # 参数通信量(GB)
            grad_size = 2e9 / (dp * 1e9)  # 梯度通信量(GB)
            comm_time = param_size / self.bandwidth + grad_size / self.bandwidth

            stage_time += compute_time + comm_time

        return stage_time

    def evaluate_strategy(self, strategy):
        """评估策略总时间"""
        stage_times = []
        for stage in strategy["stages"]:
            stage_time = self.calculate_stage_time(stage)
            stage_times.append(stage_time)

        # 总时间 = 各阶段时间之和 + (微批次数-1)*最慢阶段时间
        total_cost = sum(stage_times) + (self.num_microbatches - 1) * max(stage_times)
        return total_cost

    def compare_strategies(self, complex_path, simple_path):
        """比较两个策略"""
        complex_strategy = self.load_strategy(complex_path)
        simple_strategy = self.load_strategy(simple_path)

        complex_cost = self.evaluate_strategy(complex_strategy)
        simple_cost = self.evaluate_strategy(simple_strategy)

        improvement = (simple_cost - complex_cost) / simple_cost

        print("\n===== 策略对比结果 =====")
        print(f"Complex策略总时间: {complex_cost:.2f}秒")
        print(f"Simple策略总时间: {simple_cost:.2f}秒")
        print(f"时间提升: {improvement * 100:.1f}%")

        return {
            "complex_cost": complex_cost,
            "simple_cost": simple_cost,
            "improvement": improvement
        }


if __name__ == "__main__":
    evaluator = StrategyEvaluator(num_microbatches=32)

    # 文件路径（根据实际情况修改）
    complex_strategy_path = "/home/lthpc/nvmessd/wangjiaqian/simulator_by_alpa_layer/experiment/strategy/complex_strategy.json"
    simple_strategy_path = "/home/lthpc/nvmessd/wangjiaqian/simulator_by_alpa_layer/experiment/strategy/simple_strategy.json"

    # 执行比较
    evaluator.compare_strategies(complex_strategy_path, simple_strategy_path)